from lab_2_conf_manage_neo4j import get_neo4j_graph
from lab_1_load_split_doc import get_split_doc
from langchain_experimental.graph_transformers import LLMGraphTransformer
from langchain_community.chat_models import ChatTongyi


def construct_kg(qry):
    # Define Neo4j client
    graph = get_neo4j_graph()

    # Define LLM and LLM Graph transformer
    llm = ChatTongyi(api_key='sk-780da555f1714f4b81c1bfc553cfd74a')
    llm_transformer = LLMGraphTransformer(llm=llm)

    # Extract graph data
    graph_documents = llm_transformer.convert_to_graph_documents(get_split_doc(qry))

    # Store to neo4j
    # baseEntityLabel: True 意味着每个图文档实体都会被添加为图中的一个节点
    # include_source：True 意味着原始文档的元数据或上下文信息也可能被存储在图中
    graph.add_graph_documents(
        graph_documents,
        baseEntityLabel=True,
        include_source=True
    )

    print("Construct Knowledge Graph successful.")


if __name__ == "__main__":
    query = "Elizabeth I"
    construct_kg(query)




